import time
import os
import tkinter as tk
from PIL import ImageGrab
from pynput import keyboard
import json # <-- 用于处理 answer.json 文件

# 引入 PaddleOCR 库
try:
    # 假设这里是你的自定义 PaddleOCR 库（可能包含对结果结构的修改）
    from paddleocr import PaddleOCR
    # 定义全局变量用于存储 OCR 实例
    ocr_engine = None 
except ImportError:
    print("❌ 警告：PaddleOCR 库未找到。请确保已安装：pip install paddleocr")
    ocr_engine = None
    PaddleOCR = None 


# --- 配置 ---
SAVE_DIR = "D:\\fk2025"
ACTIVATION_KEY = keyboard.KeyCode.from_char('q') 
# 确保文件夹存在
os.makedirs(SAVE_DIR, exist_ok=True)

# --- 全局状态 ---
app = None
answer_display = None # 用于存储答案显示窗口实例
start_x, start_y = 0, 0
current_x, current_y = 0, 0
is_capturing = False

# --- 答案显示窗口类：用于在屏幕左上角显示答案 (支持多颜色) ---
class AnswerDisplayApp(tk.Toplevel):
    def __init__(self, master):
        super().__init__(master)
        
        # 窗口配置
        self.title("OCR 答案")
        self.geometry("700x350+0+0") # 尺寸 700x350
        self.overrideredirect(True)      # 移除窗口边框
        self.attributes('-topmost', True) # 置顶
        self.attributes('-alpha', 0.9)    # 半透明 (0.9 不完全透明)
        self.configure(bg='black')       # 背景色

        # 答案显示区域：使用 Text 控件支持多颜色
        self.answer_text = tk.Text(
            self, 
            bg="black", 
            fg="white", # 默认文本颜色（用于答案内容）
            insertbackground="white",
            wrap=tk.WORD, 
            font=("Arial", 10),
            padx=5, 
            pady=5
        )
        self.answer_text.pack(fill=tk.BOTH, expand=True)
        
        # 配置颜色标签 (Tags)
        # 【颜色配置】‘正确答案：’红色显示
        self.answer_text.tag_config('red_label', foreground='red', font=("Arial", 10, "bold")) 
        # 【颜色配置】查找的关键字蓝色显示
        self.answer_text.tag_config('blue_keyword', foreground='blue') 
        
        self.withdraw() # 初始隐藏窗口
        
        # 初始显示
        self.show_initial_message()

    def show_initial_message(self):
        """显示初始提示信息"""
        self.answer_text.config(state=tk.NORMAL)
        self.answer_text.delete(1.0, tk.END)
        self.answer_text.insert(tk.END, "等待截图...")
        self.answer_text.config(state=tk.DISABLED)

    def show_answer(self, keyword, display_output):
        """更新文本内容，显示关键字和答案，并应用颜色标签"""
        self.answer_text.config(state=tk.NORMAL)
        self.answer_text.delete(1.0, tk.END)
        
        # 插入查找关键字标签 (默认颜色)
        self.answer_text.insert(tk.END, "查找关键字: ")
        # 插入关键字 (蓝色)
        self.answer_text.insert(tk.END, keyword + "\n", 'blue_keyword') 
        
        # 插入正确答案标签 (红色)
        self.answer_text.insert(tk.END, "正确答案：\n", 'red_label') 
        
        # 插入答案内容 (默认白色)
        self.answer_text.insert(tk.END, display_output)
        
        self.answer_text.config(state=tk.DISABLED) # 设置为只读
        self.deiconify() # 显示窗口

    def hide_answer(self):
        """隐藏窗口并重置文本"""
        self.show_initial_message()
        self.withdraw()


# --- OCR 识别函数：使用您提供的 predict 逻辑 ---
def perform_ocr(image_path):
    """
    使用全局的 PaddleOCR 实例对图片进行识别，并提取 rec_texts。
    :param image_path: 截图文件的完整路径。
    """
    global ocr_engine 
    global answer_display # 【关键修复】将 global 声明提前到函数顶部
    
    if ocr_engine is None:
        print("❌ OCR 引擎未初始化或加载失败，跳过识别。")
        # 此时 answer_display 仍可访问，但在 ocr_engine 失败时不显示答案，直接返回。
        return

    print("\n----------------------------------------------")
    print("🧠 正在进行 OCR 识别...")
    
    try:
        # 使用 ocr.predict 方法进行推理
        results = ocr_engine.predict(input=image_path)

        if not results:
            print("🔍 未检测到可识别的文本。")
            if answer_display:
                 answer_display.show_answer("无法识别", "未检测到可识别的文本。")
            return
        
        print("✅ 识别完成！输出结果：")
        print("==============================================")
        
        recognized_texts = []
        for res in results:
            if 'rec_texts' in res:
                recognized_texts = res['rec_texts']
            elif 'res' in res and isinstance(res['res'], dict) and 'rec_texts' in res['res']:
                 recognized_texts = res['res']['rec_texts']
            else:
                print("⚠️ 警告：未在识别结果中找到 'rec_texts' 键。")

            if recognized_texts:
                print(recognized_texts[0])

        print("==============================================")
        if recognized_texts:
            qdn(recognized_texts[0])
        else:
            # 如果识别失败，则在答案区域显示信息
            if answer_display:
                answer_display.show_answer("无法识别", "未检测到可识别的文本。")
                
    except Exception as e:
        print(f"❌ OCR 识别过程中发生错误: {e}")
        # global answer_display # 【修复】此处的 global 声明已移除
        if answer_display:
             answer_display.show_answer("OCR 错误", f"识别过程中发生错误: {e}")


# --- 【更新】qdn 函数：查询 answer.json 文件并显示结果 ---
def qdn(text):
    """
    实现从当前目录的answer.json中查找答案，以text为关键字，
    输出到控制台，并将结果显示在屏幕上的答案区域。
    :param text: OCR识别结果（关键字）。
    """
    global answer_display
    
    # 确定 answer.json 的完整路径
    base_dir = os.path.dirname(os.path.abspath(__file__)) 
    json_path = os.path.join(base_dir, "answer.json")

    print("\n----------------------------------------------")
    print("🔎 正在从 answer.json 中查询答案...")
    print(f"🔑 关键字: **{text}**")
    
    if answer_display is None:
        print("⚠️ 警告：答案显示窗口未初始化。")

    if not os.path.exists(json_path):
        error_msg = f"❌ 错误：未找到文件 {json_path}。"
        print(error_msg)
        if answer_display:
            # 传入关键字和错误信息
            answer_display.show_answer(text, error_msg) 
        return

    try:
        # 读取并解析 JSON 文件
        with open(json_path, 'r', encoding='utf-8') as f:
            data = json.load(f)
            
        matches = []
        search_key = text.lower()
        
        # 查找逻辑：兼容列表和字典
        if isinstance(data, list):
            for item in data:
                item_str = str(item)
                if search_key in item_str.lower():
                    matches.append(item_str)
        
        elif isinstance(data, dict):
             for key, value in data.items():
                if search_key in key.lower() or search_key in str(value).lower():
                    matches.append(f"{key}: {value}")

        # 处理并输出结果
        if matches:
            print(f"✅ 找到 **{len(matches)}** 条匹配结果:")
            print("==============================================")
            
            # 1. 构造控制台输出文本
            console_output = '\n********\n'.join(matches)
            print(console_output) 
            print("==============================================")

            # 2. 构造屏幕显示文本，将 <br> 替换为换行符 \n。
            display_output = console_output.replace('<br>', '\n')
            
            if answer_display:
                # 调用 show_answer 时，传入关键字 text 和答案 display_output
                answer_display.show_answer(text, display_output)
        else:
            no_match_msg = "🔍 未找到匹配的答案。"
            print(no_match_msg)
            if answer_display:
                # 对于无答案的情况，传入关键字 text 和提示信息
                answer_display.show_answer(text, no_match_msg)
            
    except json.JSONDecodeError as e:
        error_msg = f"❌ 错误：answer.json 文件解析失败，请检查 JSON 格式: {e}"
        print(error_msg)
        if answer_display:
            # 对于错误，传入关键字 text 和错误信息
            answer_display.show_answer(text, error_msg)
    except Exception as e:
        error_msg = f"❌ 答案查询过程中发生未知错误: {e}"
        print(error_msg)
        if answer_display:
            # 对于错误，传入关键字 text 和错误信息
            answer_display.show_answer(text, error_msg)

    return text

# --- 截图类：负责图形绘制和截图逻辑 (此处无改动) ---
class ScreenshotApp(tk.Toplevel):
    def __init__(self, master=None):
        super().__init__(master)
        
        # 1. 窗口配置：实现绿色透明和置顶
        self.attributes('-alpha', 0.2)      
        self.attributes('-fullscreen', True) 
        self.attributes('-topmost', True)   
        self.overrideredirect(True)         
        self.withdraw()                     

        # 2. Canvas 配置：用于绘制边框
        self.canvas = tk.Canvas(self, cursor="cross", highlightthickness=0, bg='#00FF00') 
        self.canvas.pack(fill=tk.BOTH, expand=tk.YES)
        
        # 3. 鼠标事件绑定
        self.canvas.bind("<ButtonPress-1>", self.on_mouse_down)
        self.canvas.bind("<B1-Motion>", self.on_mouse_move) 
        self.canvas.bind("<ButtonRelease-1>", self.on_mouse_up)
        
        # 4. 实时坐标标签 (用于显示鼠标当前位置)
        self.coord_label = tk.Label(self.canvas, text="", bg="black", fg="yellow", font=("Arial", 10))
        self.coord_label.place(x=10, y=10)
        
        # 5. 键盘事件绑定 (按 Esc 退出截图模式)
        self.bind('<Escape>', self.cancel_capture)

    def activate_capture_mode(self):
        """激活截图模式：显示全屏透明窗口"""
        global answer_display
        if answer_display:
            answer_display.hide_answer() # 激活截图时隐藏答案窗口
        self.deiconify() 
        self.focus_force() 
        print("\n--- 截图模式已激活，请拖动鼠标选择区域或按 Esc 取消 ---")

    def cancel_capture(self, event=None):
        """取消截图操作：隐藏窗口"""
        global is_capturing
        self.withdraw() 
        is_capturing = False
        print("--- 截图模式已取消 ---")

    def on_mouse_down(self, event):
        """鼠标左键按下：记录起始点"""
        global start_x, start_y, is_capturing
        start_x, start_y = event.x_root, event.y_root
        is_capturing = True
        self.canvas.delete("rect") 
        
    def on_mouse_move(self, event):
        """鼠标左键拖动：实时绘制边框和更新坐标"""
        if is_capturing:
            global current_x, current_y
            current_x, current_y = event.x_root, event.y_root
            
            # 绘制矩形边框（亮黄色，4px）
            self.canvas.delete("rect")
            self.canvas.create_rectangle(
                start_x, start_y, current_x, current_y,
                outline="yellow", 
                width=4,          
                tags="rect"
            )
            
            # 更新鼠标位置显示
            width = abs(current_x - start_x)
            height = abs(current_y - start_y)
            self.coord_label.config(text=f"位置: ({current_x}, {current_y})\n大小: {width}x{height}")
            # 移动标签到鼠标附近
            self.coord_label.place(x=current_x + 15, y=current_y + 15)


    def on_mouse_up(self, event):
        """鼠标左键释放：执行截图并隐藏窗口"""
        global is_capturing
        
        if is_capturing:
            end_x, end_y = event.x_root, event.y_root
            saved_path = self.perform_screenshot(start_x, start_y, end_x, end_y)
            
            if saved_path:
                perform_ocr(saved_path)
        
        is_capturing = False

    def perform_screenshot(self, x1, y1, x2, y2):
        """根据坐标截图并保存文件，确保在截图前隐藏边框"""
        
        # 1. 关键步骤：在抓取屏幕内容前，瞬间隐藏包含边框的窗口
        self.withdraw() 
        
        # 2. 规范坐标
        left = min(x1, x2)
        top = min(y1, y2)
        right = max(x1, x2)
        bottom = max(y1, y2)
        
        save_path = None
        try:
            # 3. 抓取屏幕截图 (抓取的是截图框下面的实际内容)
            img = ImageGrab.grab(bbox=(left, top, right, bottom))
            
            # 构造文件名: test时间戳.jpg
            timestamp = time.strftime("%Y%m%d%H%M%S")
            file_name = f"test{timestamp}.jpg"
            save_path = os.path.join(SAVE_DIR, file_name)
            
            img.save(save_path, "JPEG")
            print(f"✅ 截图成功并保存到: {save_path}")
            
            # 返回保存的文件路径
            return save_path
            
        except Exception as e:
            print(f"❌ 截图失败: {e}")
            return None


# --- 全局键盘监听器 (用于激活截图模式) (此处无改动) ---
def on_press(key):
    """监听键盘按下事件"""
    global app, is_capturing

    # 激活截图模式的按键
    if key == ACTIVATION_KEY and not is_capturing:
        if app:
            app.activate_capture_mode()
        return

    # 按 Esc 键退出程序
    if key == keyboard.Key.esc:
        if app and app.winfo_ismapped(): # 如果截图窗口是显示的，先取消截图
            app.cancel_capture()
        else: # 否则退出整个程序
            print("\n--- 检测到 Esc 键，程序退出 ---")
            # 停止 Tkinter 主循环
            if app:
                app.master.quit()
            return False # 停止 pynput 监听器

# --- 主程序启动 (OCR 引擎初始化参数已修改) ---
if __name__ == "__main__":
    
    # 0. 初始化 OCR 引擎 (仅在程序启动时执行一次)
    if ocr_engine is None and 'PaddleOCR' in globals() and PaddleOCR is not None: 
        print("🧠 正在初始化 PaddleOCR 引擎 (首次运行将下载模型)...")
        
        try:
            ocr_engine = PaddleOCR(
                use_doc_orientation_classify=False, 
                use_doc_unwarping=False,            
                use_textline_orientation=False    
            ) 
            print("🧠 PaddleOCR 引擎初始化成功。")
        except Exception as e:
            print(f"❌ PaddleOCR 引擎初始化失败，错误信息: {e}")
        
    # 1. 启动 Tkinter 主窗口
    root = tk.Tk()
    root.withdraw() # 隐藏根窗口
    
    # 2. 实例化截图应用
    app = ScreenshotApp(root)

    # 3. 实例化答案显示窗口
    answer_display = AnswerDisplayApp(root)

    # 4. 启动 pynput 键盘监听 (在单独线程中)
    k_listener = keyboard.Listener(on_press=on_press)
    k_listener.daemon = True # 设置为守护线程
    k_listener.start()

    print("**********************************************")
    print(f"🚀 截图 OCR 助手启动，保存目录: {SAVE_DIR}")
    print(f"📌 按下键 {ACTIVATION_KEY.char.upper()} 激活截图模式")
    print("📌 激活后，按住鼠标左键拖动以绘制边框并截图")
    print("📌 截图完成后，将自动进行 OCR 识别，并在左上角显示答案")
    print("📌 随时按 Esc 键退出程序")
    print("**********************************************")
    
    # 5. 启动 Tkinter 主循环 (GUI 必须在主线程运行)
    root.mainloop()

    # 6. 清理 pynput 监听器
    k_listener.stop()
    print("程序已安全终止。")